# coding=utf-8
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import sklearn.datasets
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
# from spark_sklearn import GridSearchCV
from sklearn.svm import SVC

iris = sklearn.datasets.load_iris()
param_grid = {'C': [0.001, 0.01, 0.1, 1, 10, 100],
              'gamma': [0.001, 0.01, 0.1, 1, 10, 100]}
grid_search = GridSearchCV(SVC(), param_grid, cv=5)
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, random_state=0)

grid_search.fit(X_train, y_train)

print("Best parameters: {}".format(grid_search.best_params_))
print("Best best_estimator:\n {}".format(grid_search.best_estimator_))
print("Test set score: {:.2f}".format(grid_search.score(X_test, y_test)))

results = pd.DataFrame(grid_search.cv_results_)
# print(results.head())
scores = np.array(results.mean_test_score).reshape(6, 6)
print(scores)

sns.heatmap(scores,
            # xlabel='gamma',
            xticklabels=param_grid['gamma'],
            # ylabel='C',
            yticklabels=param_grid['C'], annot=True, fmt=".2f")
plt.show()

# 也可以这么设置参数，那么相应的核函数就在相应的参数空间内搜索内容
param_grid = [{'kernel': ['rbf'],
               'C': [0.001, 0.01, 0.1, 1, 10, 100],
               'gamma': [0.001, 0.01, 0.1, 1, 10, 100]},
              {'kernel': []}]
